from openai import OpenAI
from dotenv import load_dotenv
import os
import time
import json
import datetime
from read_data_files import read_all_txt_files

load_dotenv()

def create_client():
    """创建OpenAI客户端"""
    try:
        client = OpenAI(
            api_key = os.getenv("API_KEY"), 
            base_url = os.getenv("API_URL")
        )
        return client
    except Exception as e:
        print(f"创建客户端时出错: {e}")
        return None

def save_log(messages, response, response_time):
    """保存日志到文件"""
    try:
        # 创建时间戳作为文件名
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]  # 精确到毫秒
        filename = f"logs/llm_call_{timestamp}.json"
        
        # 准备日志数据
        log_data = {
            "timestamp": datetime.datetime.now().isoformat(),
            "request": {
                "messages": messages,
                "model": "Qwen2.5-72B-Instruct",
                "stream": True,
                "temperature": 0.7,
                "max_tokens": 2000
            },
            "response": {
                "content": response,
                "response_time_seconds": response_time
            }
        }
        
        # 保存到文件
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(log_data, f, ensure_ascii=False, indent=2)
        
        # print(f"📝 日志已保存: {filename}")
        
    except Exception as e:
        print(f"保存日志时出错: {e}")

def stream_chat(client, messages):
    """流式聊天函数"""
    try:
        print("\nAI助手正在思考...")
        print("-" * 40)
        
        start_time = time.time()
        
        response = client.chat.completions.create(
            model="Qwen2.5-72B-Instruct",
            messages=messages,
            stream=True,  # 启用流式响应
            temperature=0.7,  # 添加一些创造性
            max_tokens=2000   # 限制最大token数
        )
        
        # 收集完整回复用于历史记录
        full_response = ""
        
        print("AI助手: ", end='', flush=True)
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                content = chunk.choices[0].delta.content
                print(content, end='', flush=True)
                full_response += content
                time.sleep(0.01)  # 添加小延迟模拟打字效果
        
        end_time = time.time()
        response_time = end_time - start_time
        
        print("\n" + "-" * 40)
        
        # 保存日志
        save_log(messages, full_response, response_time)
        
        return full_response
        
    except Exception as e:
        print(f"\n调用API时出错: {e}")
        return None

def main():
    """主函数"""
    print("=" * 60)
    print("🤖 欢迎使用流式LLM聊天程序")
    print("=" * 60)
    print("提示: 输入 'quit' 或 'exit' 退出程序")
    print("提示: 输入 'clear' 清空对话历史")
    print("=" * 60)
    
    # 创建客户端
    client = create_client()
    if not client:
        print("无法创建客户端，程序退出")
        return
    
    # 初始化对话历史
    # 从 data 目录加载系统提示词
    print("📁 正在加载系统提示词...")
    data_files = read_all_txt_files()

    messages = []
    
    if data_files:
        # 将所有 txt 文件内容合并作为系统提示词
        system_content = ""
        for filename, content in data_files.items():
            system_content += f"=== {filename} ===\n{content}\n\n"
        
        print(f"✅ 成功加载 {len(data_files)} 个文件作为系统提示词")
        messages.append({"role": "system", "content": system_content.strip()})

    messages.append({"role": "system", "content": "你是山东警察学院的招生助手，你专门负责回答关于山东警察学院的招生问题。对于招生以外的问题，请礼貌地拒绝。"})
    messages.append({"role": "system", "content": "注意：回答问题的时候不要使用mardkdown格式，要使用口语化的语言。"})

    while True:
        try:
            # 获取用户输入
            user_input = input("\n👤 您: ").strip()
            
            # 检查退出命令
            if user_input.lower() in ['quit', 'exit', '退出']:
                print("\n👋 再见！感谢使用流式LLM聊天程序")
                break
            
            # 检查清空命令
            if user_input.lower() in ['clear', '清空']:
                messages = [messages[0]]  # 保留系统消息
                print("\n🗑️  对话历史已清空")
                continue
            
            # 检查空输入
            if not user_input:
                print("请输入有效的问题")
                continue
            
            # 添加用户消息到历史
            messages.append({"role": "user", "content": user_input})
            
            # 流式调用
            response = stream_chat(client, messages)
            
            if response:
                # 添加助手回复到历史
                messages.append({"role": "assistant", "content": response})
                
                # 限制历史长度（保留最近10轮对话）
                if len(messages) > 21:  # 1个系统消息 + 20个对话消息
                    messages = [messages[0]] + messages[-20:]
            else:
                print("获取回复失败，请重试")
                # 移除失败的用户消息
                messages.pop()
        
        except KeyboardInterrupt:
            print("\n\n👋 程序被用户中断，再见！")
            break
        except Exception as e:
            print(f"\n程序运行时出错: {e}")
            print("请重试...")

if __name__ == "__main__":
    main() 